# isomorphic_noise_generator.py
import numpy as np
import torch
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import hashlib

class IsomorphicNoiseGenerator:
    """
    Generates isomorphic Gaussian noise using AES-256 encryption
    with shared seeds for reproducible noise patterns across clients.
    
    Args:
        seed (bytes): Shared secret seed
        dimensions (int): Dimension of noise vector
        sigma (float): Standard deviation of Gaussian noise (default=1.0)
    """
    
    def __init__(self, seed, dimensions, sigma=1.0):
        self.seed = seed
        self.dim = dimensions
        self.sigma = sigma
        self.aes_key = self._derive_aes_key(seed)
    
    def _derive_aes_key(self, seed):
        """Derive 256-bit AES key from seed using SHA-256"""
        return hashlib.sha256(seed).digest()
    
    def generate_noise(self, timestamp):
        """
        Generate isomorphic noise for given timestamp
        
        Args:
            timestamp (int): Training round identifier
            
        Returns:
            torch.Tensor: Noise tensor of shape (self.dim)
        """
        # AES encryption with timestamp
        nonce = timestamp.to_bytes(16, 'big')
        cipher = Cipher(algorithms.AES(self.aes_key), modes.CTR(nonce), backend=default_backend())
        encryptor = cipher.encryptor()
        
        # Generate random bytes and map to Gaussian distribution
        random_bytes = encryptor.update(b'\x00' * self.dim * 4) + encryptor.finalize()
        random_floats = np.frombuffer(random_bytes, dtype=np.float32)
        gaussian_noise = random_floats * self.sigma
        return torch.tensor(gaussian_noise[:self.dim])
